#pragma once

#include <NvInfer.h>
#include <NvInferRuntimeCommon.h>

#include <memory>
#include <opencv2/opencv.hpp>
#include <string>
#include <vector>

#include "nn_common.hpp"

using namespace nvinfer1;

class TRTDetector {
 private:
  std::string engine_path_;

  TRT::TRTLogger logger_;

  UniquePtr<ICudaEngine> engine_;
  UniquePtr<IExecutionContext> context_;

  float conf_thresh_, nms_thresh_;
  std::vector<void *> bindings_;
  std::vector<size_t> bingings_size_;
  int idx_in_, idx_out_;

  // float data_[24* 3 * 640 * 640];
  float *data_;
  cudaStream_t stream_;

  void PreProcess(cv::Mat raw);
  std::vector<TRT::Detection> PostProcess(std::vector<float> prob, int width,
                                          int height);

  bool CreateContex();
  bool InitMemory();

 public:
  TRTDetector();
  TRTDetector(const std::string &engine_file_path,
              float conf_thresh = Yolo::CONF_THRESH,
              float nms_thresh = Yolo::NMS_THRESH);
  ~TRTDetector();

  bool LoadEngine(UniquePtr<ICudaEngine> engine);
  void SetEnginePath(const std::string &engine_file_path);
  void Init();

  std::vector<TRT::Detection> Infer(const cv::Mat &input);

  bool TestInfer(const std::string &file_name);
};
